from pydantic import BaseModel, Field, field_validator
from typing import Optional, Literal, List, Dict, Any
from enum import Enum

# ============================================================================
# 通用枚举定义
# ============================================================================

class AspectRatio(str, Enum):
    """宽高比枚举"""
    RATIO_1_1 = "1:1"      # 1328x1328 正方形
    RATIO_4_3 = "4:3"      # 1472x1104 传统
    RATIO_3_2 = "3:2"      # 1584x1056 摄影
    RATIO_16_9 = "16:9"    # 1664x936  宽屏
    RATIO_21_9 = "21:9"    # 2016x864  超宽屏
    RATIO_3_4 = "3:4"      # 竖屏
    RATIO_9_16 = "9:16"    # 竖屏视频

class Style(str, Enum):
    """生成风格枚举"""
    REALISTIC = "realistic"
    ARTISTIC = "artistic" 
    ANIME = "anime"

class VideoResolution(str, Enum):
    """视频分辨率枚举"""
    HD_720P = "720p"
    FHD_1080P = "1080p"

# ============================================================================
# 图片生成Schema
# ============================================================================

class ImageGenerateRequest(BaseModel):
    """
    图片生成请求 - 支持2种模式
    1. 纯文本生图：只填prompt和aspect_ratio等基础参数
    2. 模板生图：提供template_image_url + user_image_url，将用户图片与模板融合
    """
    prompt: str = Field(..., min_length=1, max_length=800, description="图片生成提示词")
    
    # 基础参数
    aspect_ratio: Optional[AspectRatio] = Field(default=AspectRatio.RATIO_16_9, description="图片宽高比")
    style: Optional[Style] = Field(default=Style.REALISTIC, description="生成风格")
    seed: Optional[int] = Field(default=-1, description="随机种子，-1为随机")
    
    # 模板生图参数 - 需要两张图片
    template_image_url: Optional[str] = Field(default=None, description="模板图片URL")
    template_image_base64: Optional[str] = Field(default=None, description="模板图片Base64")
    user_image_url: Optional[str] = Field(default=None, description="用户图片URL，与模板融合")
    user_image_base64: Optional[str] = Field(default=None, description="用户图片Base64")
    fusion_strength: Optional[float] = Field(default=0.6, ge=0.1, le=1.0, description="图片融合强度")
    
    @field_validator('template_image_url', 'user_image_url')
    @classmethod
    def validate_template_mode(cls, v, info):
        # 如果提供了template_image_url，则user_image_url也必须提供
        template_provided = bool(info.data.get('template_image_url') or info.data.get('template_image_base64'))
        user_provided = bool(info.data.get('user_image_url') or info.data.get('user_image_base64'))
        
        if template_provided and not user_provided:
            raise ValueError('模板生图模式需要同时提供template_image_url和user_image_url')
        if user_provided and not template_provided:
            raise ValueError('模板生图模式需要同时提供template_image_url和user_image_url')
        
        return v

class ImageGenerateResponse(BaseModel):
    """图片生成响应"""
    success: bool
    task_id: Optional[str] = None
    image_url: Optional[str] = None
    cost: Optional[float] = None
    error_message: Optional[str] = None
    
    # 任务状态
    status: Optional[str] = None  # "queued", "processing", "completed", "failed"
    message: Optional[str] = None
    
    # 请求信息回显
    aspect_ratio: Optional[str] = None

# ============================================================================
# 视频生成Schema
# ============================================================================

class VideoGenerateRequest(BaseModel):
    """
    视频生成请求 - 支持2种模式
    1. 纯文本生视频：只填prompt和duration等基础参数
    2. 图片转视频：提供source_image_url (通常是图片生成接口的输出)
    """
    prompt: str = Field(..., min_length=1, max_length=800, description="视频生成提示词")
    
    # 基础参数
    duration: int = Field(default=5, description="视频时长（秒）5或10")
    aspect_ratio: Optional[AspectRatio] = Field(default=AspectRatio.RATIO_16_9, description="视频宽高比")
    resolution: Optional[VideoResolution] = Field(default=VideoResolution.HD_720P, description="视频分辨率")
    style: Optional[Style] = Field(default=Style.REALISTIC, description="视频风格")
    seed: Optional[int] = Field(default=-1, description="随机种子，-1为随机")
    
    # 图片转视频参数
    source_image_url: Optional[str] = Field(default=None, description="源图片URL，提供时为图片转视频模式")
    source_image_base64: Optional[str] = Field(default=None, description="源图片Base64")
    image_strength: Optional[float] = Field(default=0.8, ge=0.1, le=1.0, description="图片影响视频的强度")
    
    @field_validator('duration')
    @classmethod
    def validate_duration(cls, v):
        if v not in [5, 10]:
            raise ValueError('duration must be 5 or 10')
        return v

class VideoGenerateResponse(BaseModel):
    """视频生成响应"""
    success: bool
    task_id: Optional[str] = None
    video_url: Optional[str] = None
    cost: Optional[float] = None
    estimated_time: Optional[int] = None
    error_message: Optional[str] = None
    
    # 任务状态
    status: Optional[str] = None  # "queued", "processing", "completed", "failed"
    message: Optional[str] = None
    progress: Optional[int] = None  # 0-100
    
    # 请求信息回显
    generation_type: Optional[str] = None  # "text_to_video" or "image_to_video"
    aspect_ratio: Optional[str] = None
    duration: Optional[int] = None

# ============================================================================
# 组合流程Schema - 模板生视频一站式接口
# ============================================================================

class TemplateToVideoRequest(BaseModel):
    """
    模板生视频请求 - 一站式接口
    用户提供模板图片+普通图片+提示词，自动完成：
    1. 图片融合 (template + user image)
    2. 图片转视频
    """
    prompt: str = Field(..., min_length=1, max_length=800, description="视频内容描述")
    
    # 必需的图片参数
    template_image_url: str = Field(..., description="模板图片URL")
    user_image_url: str = Field(..., description="用户图片URL")
    
    # 可选的图片参数
    template_image_base64: Optional[str] = Field(default=None, description="模板图片Base64")
    user_image_base64: Optional[str] = Field(default=None, description="用户图片Base64")
    
    # 第一步：图片融合参数
    image_prompt: Optional[str] = Field(default=None, description="图片融合提示词，不填则使用主prompt")
    fusion_strength: Optional[float] = Field(default=0.6, ge=0.1, le=1.0, description="图片融合强度")
    
    # 第二步：视频生成参数
    duration: int = Field(default=5, description="视频时长（秒）")
    aspect_ratio: Optional[AspectRatio] = Field(default=AspectRatio.RATIO_16_9, description="视频宽高比")
    resolution: Optional[VideoResolution] = Field(default=VideoResolution.HD_720P, description="视频分辨率")
    style: Optional[Style] = Field(default=Style.REALISTIC, description="视频风格")
    image_strength: Optional[float] = Field(default=0.8, ge=0.1, le=1.0, description="图片转视频时的图片影响强度")
    
    # 通用参数
    seed: Optional[int] = Field(default=-1, description="随机种子")
    
    @field_validator('duration')
    @classmethod
    def validate_duration(cls, v):
        if v not in [5, 10]:
            raise ValueError('duration must be 5 or 10')
        return v

class TemplateToVideoResponse(BaseModel):
    """模板生视频响应"""
    success: bool
    task_id: Optional[str] = None
    
    # 最终结果
    video_url: Optional[str] = None
    intermediate_image_url: Optional[str] = None  # 中间生成的融合图片URL
    
    # 成本和时间
    total_cost: Optional[float] = None
    image_cost: Optional[float] = None  # 图片融合成本
    video_cost: Optional[float] = None  # 视频生成成本
    estimated_time: Optional[int] = None
    
    # 错误处理
    error_message: Optional[str] = None
    step_failed: Optional[str] = None  # 失败的步骤："image_fusion" or "video_generation"
    
    # 任务状态
    status: Optional[str] = None  # "queued", "image_processing", "video_processing", "completed", "failed"
    message: Optional[str] = None
    progress: Optional[int] = None  # 0-100
    current_step: Optional[str] = None  # "图片融合中", "视频生成中"
    
    # 请求信息回显
    aspect_ratio: Optional[str] = None
    duration: Optional[int] = None

# ============================================================================
# 任务状态查询Schema
# ============================================================================

class TaskStatusResponse(BaseModel):
    """通用任务状态响应"""
    task_id: str
    status: str  # "queued", "processing", "completed", "failed"
    progress: Optional[int] = None  # 0-100
    message: Optional[str] = None
    
    # 完成时的结果
    result_url: Optional[str] = None  # 图片或视频URL
    cost: Optional[float] = None
    
    # 错误信息
    error: Optional[str] = None
    
    # 时间信息
    created_at: Optional[str] = None
    completed_at: Optional[str] = None
    processing_time: Optional[int] = None  # 处理时间（秒）